"""Model evaluate.

Use Coco2017 validation dataset, evaluate the precision of PyTorch model or OM
model by computing mAP metric.

1. Evaluate the precision of OM model.
Example:
>>> python tools/evaluate.py \
>>>     mmdetection/configs/yolox/yolox_s_8xb8-300e_coco.py \
>>>     --mode offline \
>>>     --om_model models/yolox_s_bs1.om

2. Evaluate the precision of PyTorch model.
Example:
>>> python tools/evaluate.py \
>>>     mmdetection/configs/yolox/yolox_s_8xb8-300e_coco.py \
>>>     --mode online \
>>>     --checkpoint models/yolox_s_8x8_300e_coco_20211121_095711-4592a793.pth \
>>>     --batch_size 1

3. Evaluate the precision of PyTorch model, but use the offline data consistent
   with the OM model. It is very helpful to locate accuracy problem.
Example:
>>> python tools/evaluate.py \
>>>     mmdetection/configs/yolox/yolox_s_8xb8-300e_coco.py \
>>>     --mode online \
>>>     --checkpoint models/yolox_s_8x8_300e_coco_20211121_095711-4592a793.pth \
>>>     --batch_size 1 \
>>>     --use_offline_data \
>>>     --input_shape 640 640
"""

import argparse
from pathlib import Path
from npu_mmdet.apis import offline_evaluate, online_evaluate, \
                           online_evaluate_with_offline_data


def main():
    parser = argparse.ArgumentParser(
                        description='evaluate om model.')
    parser.add_argument('config', type=str, help='path to model config')
    parser.add_argument('--mode', type=str,
                        default='offline', choices=['online', 'offline'],
                        help='mode of evaluation.')
    parser.add_argument('--om_model', type=str, help='path to OM model.')
    parser.add_argument('--checkpoint', type=str,
                        help='path to PyTorch checkpoint file.')
    parser.add_argument('--batch_size', type=int,
                        default=1,
                        help='batch size of model inputs')
    parser.add_argument('--input_shape', type=int, nargs=2,
                        help='model input shape, format: [H, W]')
    parser.add_argument('--use_offline_data', action='store_true',
                        help='when `mode` equal to "online", whether use '
                             'offline data to evaluate?')
    args = parser.parse_args()

    if args.mode == 'offline':
        assert args.om_model and Path(args.om_model).is_file(), \
            f'The value of argument `--om_model` is invalid: {args.om_model}'
        offline_evaluate(args.config, args.om_model)
    else:
        assert args.checkpoint and Path(args.checkpoint).is_file(), \
            'The value of argument `--checkpoint` is invalid: ' \
            f'{args.checkpoint}'
        if not args.use_offline_data:
            online_evaluate(args.config, args.checkpoint,
            batch_size=args.batch_size)
        else:
            assert args.input_shape and all(x > 0 for x in args.input_shape), \
                'The value of argument `--input_shape` is invalid: ' \
                f'{args.input_shape}'
            online_evaluate_with_offline_data(
                args.config, args.checkpoint,
                batch_size=args.batch_size, input_shape=args.input_shape)


if __name__ == '__main__':
    main()
